use std::env;
use std::fs::{self, File};
use std::io::{self, Read, Write, BufReader, BufWriter, Seek, SeekFrom};
use std::path::Path;
use rand::Rng;

const CHUNK_SIZE: usize = 1024 * 1024; // 1MB 的块大小

fn split_file(input_path: &str, num_parts: usize) -> io::Result<()> {
    let file = File::open(input_path)?;
    let file_size = file.metadata()?.len();
    
    // 计算基础分片大小
    let base_chunk_size = (file_size as usize + num_parts - 1) / num_parts;
    let mut reader = BufReader::new(file);
    let mut rng = rand::thread_rng();
    
    // 计算所有分片的大小
    let mut chunk_sizes = Vec::with_capacity(num_parts);
    let mut remaining_size = file_size as usize;
    
    for i in 0..num_parts {
        if i == num_parts - 1 {
            // 最后一个分片使用剩余的所有大小
            chunk_sizes.push(remaining_size);
        } else {
            // 随机浮动范围在基础大小的 ±10% 之间
            let variation = (base_chunk_size as f64 * 0.1) as usize;
            let random_variation = rng.gen_range(0..=variation * 2) as i64 - variation as i64;
            let chunk_size = (base_chunk_size as i64 + random_variation).max(1) as usize;
            
            // 确保不会超出剩余大小
            let actual_size = chunk_size.min(remaining_size);
            chunk_sizes.push(actual_size);
            remaining_size -= actual_size;
        }
    }
    
    // 打乱分片大小的顺序，使大小分布更随机
    for i in 0..num_parts-1 {
        let j = rng.gen_range(i..num_parts);
        chunk_sizes.swap(i, j);
    }
    
    // 创建对应大小的缓冲区
    let max_chunk_size = *chunk_sizes.iter().max().unwrap();
    let mut buffer = vec![0; max_chunk_size];
    
    // 写入分片
    for (chunk_number, &chunk_size) in chunk_sizes.iter().enumerate() {
        let bytes_read = reader.read(&mut buffer[..chunk_size])?;
        if bytes_read == 0 {
            break;
        }
        
        let chunk_path = format!("{}.part{}", input_path, chunk_number);
        let mut chunk_file = BufWriter::new(File::create(&chunk_path)?);
        chunk_file.write_all(&buffer[..bytes_read])?;
        chunk_file.flush()?;
        
        println!("分片 {} 大小: {} 字节", chunk_number, bytes_read);
    }
    
    println!("文件已分割成 {} 个部分，基础大小约 {} 字节", num_parts, base_chunk_size);
    Ok(())
}

fn merge_files(input_pattern: &str, output_path: &str) -> io::Result<()> {
    let mut output_file = BufWriter::new(File::create(output_path)?);
    let mut chunk_number = 0;
    let mut total_bytes = 0u64;
    
    loop {
        let chunk_path = format!("{}.part{}", input_pattern, chunk_number);
        if !Path::new(&chunk_path).exists() {
            break;
        }
        
        let chunk_file = File::open(&chunk_path)?;
        let mut chunk_reader = BufReader::new(chunk_file);
        
        let bytes_copied = io::copy(&mut chunk_reader, &mut output_file)?;
        total_bytes += bytes_copied;
        chunk_number += 1;
    }
    
    output_file.flush()?;
    
    if chunk_number == 0 {
        return Err(io::Error::new(
            io::ErrorKind::NotFound,
            "未找到任何分片文件"
        ));
    }
    
    println!("文件已合并完成，总大小: {} 字节", total_bytes);
    Ok(())
}

fn main() {
    let args: Vec<String> = env::args().collect();
    
    if args.len() < 3 {
        println!("使用方法:");
        println!("分割文件: {} split <文件路径> <分片数量>", args[0]);
        println!("合并文件: {} merge <文件路径> <输出文件路径>", args[0]);
        return;
    }
    
    match args[1].as_str() {
        "split" => {
            if args.len() < 4 {
                println!("分割文件需要指定分片数量");
                return;
            }
            let num_parts = match args[3].parse::<usize>() {
                Ok(n) if n > 0 => n,
                _ => {
                    println!("分片数量必须是大于0的整数");
                    return;
                }
            };
            if let Err(e) = split_file(&args[2], num_parts) {
                eprintln!("分割文件时出错: {}", e);
            }
        }
        "merge" => {
            if args.len() < 4 {
                println!("合并文件需要指定输出文件路径");
                return;
            }
            if let Err(e) = merge_files(&args[2], &args[3]) {
                eprintln!("合并文件时出错: {}", e);
            }
        }
        _ => {
            println!("未知命令: {}", args[1]);
        }
    }
}
